#-*- coding: utf-8 -*-
# 这是一个示例 Python 脚本。

# 按 Shift+F10 执行或将其替换为您的代码。
# 按 双击 Shift 在所有地方搜索类、文件、工具窗口、操作和设置。

from whisper import *
from setting import *

from recorder import Recorder
from chatgpt import *
from whisper import *

from PyQt5.Qt import (
    QApplication,
    QThread,
    pyqtSignal,
    QMainWindow,
    QWidget,
    QLabel,
    QPushButton,
    QVBoxLayout,
    QHBoxLayout,
    QTextEdit,
    QStatusBar
)

import sys

app = QApplication(sys.argv)

# 音频文件默认处理目录
audio_dir = "D:\\MyCode\\py-learning\\ai\\data\\"


class TransWorker(QThread):
    trans_completed_signal = pyqtSignal(str)

    def __init__(self, file: str):
        super().__init__()
        self.file = file

    def run(self):
        print("开始翻译...{0}".format(self.file))
        trans_result: str = trans3(self.file)
        print("翻译结束...{0}".format(self.file))
        self.trans_completed_signal.emit(trans_result)


class ChatGPTWorker(QThread):
    chatgpt_completed_signal = pyqtSignal(str)

    def __init__(self, prompt: str):
        super().__init__()
        self.prompt = prompt

    def run(self):
        print("开始生成病历...")
        chatgpt_result = chat35(self.prompt)
        print("病历生成结束...")
        self.chatgpt_completed_signal.emit(chatgpt_result)


class MainWindow(QMainWindow):

    def __init__(self, rec: Recorder):
        super(MainWindow, self).__init__()

        self.txt_audio_record = None
        self.label1 = None
        self.txt_medical_record = None
        self.curr_mp3_file = None
        self.trans_text = None
        self.statusbar = None
        self.chatgpt_worker = None
        self.trans_worker = None
        self.recorder = rec
        self.btnPlay = QPushButton(self)
        self.btnStartRecord = QPushButton(self)
        self.btnEndRecord = QPushButton(self)
        self.btnReset = QPushButton(self)
        self.btnGenerateMedicalRecord = QPushButton(self)

        self.setup_main_window(rec)

    def setup_main_window(self, rec):
        widget = QWidget()

        self.setWindowTitle("大模型语音模型DEMO")
        self.resize(800, 800)
        btn_layout = QHBoxLayout()
        page_layout = QVBoxLayout()

        page_layout.addLayout(btn_layout)

        self.btnPlay.setText("播放并查看波形")
        self.btnPlay.clicked.connect(rec.play)
        btn_layout.addWidget(self.btnPlay)

        self.btnStartRecord.setText("开始录音")
        self.btnStartRecord.clicked.connect(self.start_record)
        btn_layout.addWidget(self.btnStartRecord)

        self.btnEndRecord.setText("结束录音")
        self.btnEndRecord.clicked.connect(self.stop_record)
        btn_layout.addWidget(self.btnEndRecord)

        self.btnGenerateMedicalRecord.setText("生成病历")
        self.btnGenerateMedicalRecord.clicked.connect(self.chatgpt)
        btn_layout.addWidget(self.btnGenerateMedicalRecord)

        self.btnReset.setText("重置")
        self.btnReset.clicked.connect(self.reset)
        btn_layout.addWidget(self.btnReset)

        self.label1 = QLabel(self)
        self.label1.setText("mp3文件路径~~")
        page_layout.addWidget(self.label1)
        rec.generated_mp3.connect(self.on_generated_mp3)

        self.txt_audio_record = QTextEdit()
        self.txt_audio_record.setText("录音后，自动生成语音文本。")
        page_layout.addWidget(self.txt_audio_record)
        page_layout.setStretchFactor(self.txt_audio_record, 2)

        self.txt_medical_record = QTextEdit()
        self.txt_medical_record.setText("录音后，请点击\"生病病历\"按钮，调用大模型生成病人病历")

        page_layout.addWidget(self.txt_medical_record)
        page_layout.setStretchFactor(self.txt_medical_record, 10)

        widget.setLayout(page_layout)
        self.setCentralWidget(widget)

        self.statusbar = QStatusBar(self)
        self.statusbar.setToolTip("状态栏")
        self.setStatusBar(self.statusbar)
        self.statusbar.showMessage("状态栏")

        self.reset()

    def chatgpt(self):
        self.txt_medical_record.setText("正在生成病历，请稍候...")

        prompt: str = "你的角色是一个医生，请帮助生成一份病历。患者基本信息：姓名: 王五 年龄: 60岁 性别: 男。"

        self.chatgpt_worker = ChatGPTWorker(prompt)
        self.chatgpt_worker.chatgpt_completed_signal.connect(self.on_chatgpt_completed)
        self.chatgpt_worker.start()

    def trans(self):
        self.trans_worker = TransWorker(self.curr_mp3_file)
        self.trans_worker.trans_completed_signal.connect(self.on_trans_completed)
        self.trans_worker.start()

    def reset(self):
        self.label1.setText("当前无可用音频文件")

        self.btnPlay.setEnabled(False)
        self.btnEndRecord.setEnabled(False)
        self.btnGenerateMedicalRecord.setEnabled(False)

        self.btnStartRecord.setEnabled(True)

    def start_record(self):
        self.recorder.start()
        self.btnStartRecord.setEnabled(False)
        self.btnEndRecord.setEnabled(True)

    def stop_record(self):
        self.recorder.save()
        self.btnPlay.setEnabled(True)
        self.btnGenerateMedicalRecord.setEnabled(True)
        self.btnEndRecord.setEnabled(False)

    def on_generated_mp3(self, file):
        print("on_generated_mp3: " + file)

        self.label1.setText(file)
        self.curr_mp3_file = file
        self.txt_audio_record.setText("正在语音转文字，请稍候...")

        self.trans()

    def on_chatgpt_completed(self, result):
        self.txt_medical_record.setText(result)

    def on_trans_completed(self, result):
        self.trans_text = result
        self.txt_audio_record.setText(result)
        self.btnGenerateMedicalRecord.setEnabled(True)


# 按间距中的绿色按钮以运行脚本。
if __name__ == '__main__':
    init()
    recorder = Recorder(data_dir=audio_dir)
    window = MainWindow(recorder)
    window.show()
    # 等待窗口停止
    sys.exit(app.exec())
